{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ff4be5f3",
   "metadata": {},
   "source": [
    "# VectorStore-Backed Memory\n",
    "\n",
    "`VectorStoreRetrieverMemory` stores memories in a VectorDB and queries the top-K most \"salient\" docs every time it is called.\n",
    "\n",
    "This differs from most of the other Memory classes in that it doesn't explicitly track the order of interactions.\n",
    "\n",
    "In this case, the \"docs\" are previous conversation snippets. This can be useful to refer to relevant pieces of information that the AI was told earlier in the conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "da3384db",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "from langchain.embeddings.openai import OpenAIEmbeddings\n",
    "from langchain.llms import OpenAI\n",
    "from langchain.memory import VectorStoreRetrieverMemory\n",
    "from langchain.chains import ConversationChain\n",
    "from langchain.prompts import PromptTemplate"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c2e7abdf",
   "metadata": {},
   "source": [
    "### Initialize your VectorStore\n",
    "\n",
    "Depending on the store you choose, this step may look different. Consult the relevant VectorStore documentation for more details."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "eef56f65",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import faiss\n",
    "\n",
    "from langchain.docstore import InMemoryDocstore\n",
    "from langchain.vectorstores import FAISS\n",
    "\n",
    "\n",
    "embedding_size = 1536 # Dimensions of the OpenAIEmbeddings\n",
    "index = faiss.IndexFlatL2(embedding_size)\n",
    "embedding_fn = OpenAIEmbeddings().embed_query\n",
    "vectorstore = FAISS(embedding_fn, index, InMemoryDocstore({}), {})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f4bdf92",
   "metadata": {},
   "source": [
    "### Create your the VectorStoreRetrieverMemory\n",
    "\n",
    "The memory object is instantiated from any VectorStoreRetriever."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "e00d4938",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# In actual usage, you would set `k` to be a higher value, but we use k=1 to show that\n",
    "# the vector lookup still returns the semantically relevant information\n",
    "retriever = vectorstore.as_retriever(search_kwargs=dict(k=1))\n",
    "memory = VectorStoreRetrieverMemory(retriever=retriever)\n",
    "\n",
    "# When added to an agent, the memory object can save pertinent information from conversations or used tools\n",
    "memory.save_context({\"input\": \"My favorite food is pizza\"}, {\"output\": \"thats good to know\"})\n",
    "memory.save_context({\"input\": \"My favorite sport is soccer\"}, {\"output\": \"...\"})\n",
    "memory.save_context({\"input\": \"I don't the Celtics\"}, {\"output\": \"ok\"}) # "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2fe28a28",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input: My favorite sport is soccer\n",
      "output: ...\n"
     ]
    }
   ],
   "source": [
    "# Notice the first result returned is the memory pertaining to tax help, which the language model deems more semantically relevant\n",
    "# to a 1099 than the other documents, despite them both containing numbers.\n",
    "print(memory.load_memory_variables({\"prompt\": \"what sport should i watch?\"})[\"history\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6d2569f",
   "metadata": {},
   "source": [
    "## Using in a chain\n",
    "Let's walk through an example, again setting `verbose=True` so we can see the prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "ebd68c10",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: My favorite food is pizza\n",
      "output: thats good to know\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: Hi, my name is Perry, what's up?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "\" Hi Perry, I'm doing well. How about you?\""
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm = OpenAI(temperature=0) # Can be any valid LLM\n",
    "_DEFAULT_TEMPLATE = \"\"\"The following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
    "\n",
    "Relevant pieces of previous conversation:\n",
    "{history}\n",
    "\n",
    "(You do not need to use these pieces of information if not relevant)\n",
    "\n",
    "Current conversation:\n",
    "Human: {input}\n",
    "AI:\"\"\"\n",
    "PROMPT = PromptTemplate(\n",
    "    input_variables=[\"history\", \"input\"], template=_DEFAULT_TEMPLATE\n",
    ")\n",
    "conversation_with_summary = ConversationChain(\n",
    "    llm=llm, \n",
    "    prompt=PROMPT,\n",
    "    # We set a very low max_token_limit for the purposes of testing.\n",
    "    memory=memory,\n",
    "    verbose=True\n",
    ")\n",
    "conversation_with_summary.predict(input=\"Hi, my name is Perry, what's up?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "86207a61",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: My favorite sport is soccer\n",
      "output: ...\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: what's my favorite sport?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' You told me earlier that your favorite sport is soccer.'"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Here, the basketball related content is surfaced\n",
    "conversation_with_summary.predict(input=\"what's my favorite sport?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "8c669db1",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: My favorite food is pizza\n",
      "output: thats good to know\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: Whats my favorite food\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' You said your favorite food is pizza.'"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Even though the language model is stateless, since relavent memory is fetched, it can \"reason\" about the time.\n",
    "# Timestamping memories and data is useful in general to let the agent determine temporal relevance\n",
    "conversation_with_summary.predict(input=\"Whats my favorite food\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "8c09a239",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "\u001b[1m> Entering new ConversationChain chain...\u001b[0m\n",
      "Prompt after formatting:\n",
      "\u001b[32;1m\u001b[1;3mThe following is a friendly conversation between a human and an AI. The AI is talkative and provides lots of specific details from its context. If the AI does not know the answer to a question, it truthfully says it does not know.\n",
      "\n",
      "Relevant pieces of previous conversation:\n",
      "input: Hi, my name is Perry, what's up?\n",
      "response:  Hi Perry, I'm doing well. How about you?\n",
      "\n",
      "(You do not need to use these pieces of information if not relevant)\n",
      "\n",
      "Current conversation:\n",
      "Human: What's my name?\n",
      "AI:\u001b[0m\n",
      "\n",
      "\u001b[1m> Finished chain.\u001b[0m\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "' Your name is Perry.'"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# The memories from the conversation are automatically stored,\n",
    "# since this query best matches the introduction chat above,\n",
    "# the agent is able to 'remember' the user's name.\n",
    "conversation_with_summary.predict(input=\"What's my name?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df27c7dc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
